#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import requests
import math
import time
import random
import logging
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('TileDownloader')

# 定义备选瓦片源 - 进一步优化海域瓦片源选择
TILE_SOURCES = {
    "osm": "https://{s}.tile.openstreetmap.org/{z}/{x}/{y}.png",
    "otm": "https://{s}.tile.opentopomap.org/{z}/{x}/{y}.png",
    "esri": "https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}",
    "esri_ocean": "https://server.arcgisonline.com/ArcGIS/rest/services/Ocean/World_Ocean_Base/MapServer/tile/{z}/{y}/{x}", # 海洋底图，对海域更友好
    "carto": "https://cartodb-basemaps-{s}.global.ssl.fastly.net/light_all/{z}/{x}/{y}.png",
    "google": "https://mt{s}.google.com/vt/lyrs=m&x={x}&y={y}&z={z}",
    "google_satellite": "https://mt{s}.google.com/vt/lyrs=s&x={x}&y={y}&z={z}", # 谷歌卫星图，对海域有更好支持
    "google_terrain": "https://mt{s}.google.com/vt/lyrs=p&x={x}&y={y}&z={z}", # 谷歌地形图
    "google_hybrid": "https://mt{s}.google.com/vt/lyrs=y&x={x}&y={y}&z={z}"  # 谷歌混合图
}

# 每个瓦片源对应的子域
SUBDOMAINS = {
    "osm": ["a", "b", "c"],
    "otm": ["a", "b", "c"],
    "esri": [],  # ESRI不使用子域
    "esri_ocean": [], # ESRI海洋底图不使用子域
    "carto": ["a", "b", "c", "d"],
    "google": ["0", "1", "2", "3"],
    "google_satellite": ["0", "1", "2", "3"]
}

# 海域瓦片优先级排序 - 针对海域下载优先使用这些源
OCEAN_TILE_PRIORITY = ["esri_ocean", "google_satellite", "google_hybrid", "esri", "google", "google_terrain", "carto", "osm", "otm"]

class TileDownloader:
    """
    瓦片地图下载器，用于下载指定区域的地图瓦片
    增强版本支持多瓦片源自动切换和失败重试
    """
    
    def __init__(self, tile_url_template=None, output_dir="./tiles", 
                 max_workers=3, retries=5, timeout=15, source_type="osm"):
        """
        初始化下载器
        
        参数:
            tile_url_template: 瓦片URL模板，使用{z}、{x}、{y}、{s}作为占位符
            output_dir: 输出目录
            max_workers: 最大并行下载线程数
            retries: 下载失败重试次数
            timeout: 请求超时时间（秒）
            source_type: 瓦片源类型，可选值: osm, otm, esri, esri_ocean, carto, google, google_satellite
        """
        self.source_type = source_type
        
        # 检测是否可能是海域地图，如果是推荐使用海域专用源
        self.is_ocean_map = False
        
        # 如果未提供URL模板，使用默认瓦片源
        if tile_url_template is None:
            if source_type in TILE_SOURCES:
                tile_url_template = TILE_SOURCES[source_type]
                logger.info(f"使用预设瓦片源: {source_type}")
            else:
                tile_url_template = TILE_SOURCES["osm"]  # 默认使用OSM
                logger.warning(f"未知瓦片源类型: {source_type}，使用默认OSM瓦片源")
        
        self.tile_url_template = tile_url_template
        self.output_dir = output_dir
        self.max_workers = max_workers
        self.retries = retries
        self.timeout = timeout
        self.session = requests.Session()
        
        # 获取当前瓦片源的子域列表
        self.subdomains = SUBDOMAINS.get(source_type, ["a"])
        
        # 统计信息
        self.total_requests = 0
        self.successful_requests = 0
        self.failed_requests = 0
        self.source_success_rates = {}  # 跟踪每个源的成功率
        
        # 更新请求头，模拟真实浏览器
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Accept': 'image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
            'Referer': 'https://www.openstreetmap.org/',
            'Connection': 'keep-alive',
            'Cache-Control': 'no-cache',
            'Pragma': 'no-cache'
        })
        
        # 确保输出目录存在
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
    
    def deg2num(self, lat_deg, lon_deg, zoom):
        """
        将经纬度转换为瓦片编号
        
        参数:
            lat_deg: 纬度（度）
            lon_deg: 经度（度）
            zoom: 缩放级别
            
        返回:
            (x, y): 瓦片坐标
        """
        lat_rad = math.radians(lat_deg)
        n = 2.0 ** zoom
        x = int((lon_deg + 180.0) / 360.0 * n)
        y = int((1.0 - math.log(math.tan(lat_rad) + (1 / math.cos(lat_rad))) / math.pi) / 2.0 * n)
        return (x, y)
    
    def format_tile_url(self, zoom, x, y):
        """
        格式化瓦片URL，替换占位符
        
        参数:
            zoom: 缩放级别
            x: 瓦片X坐标
            y: 瓦片Y坐标
            
        返回:
            格式化后的URL
        """
        # 随机选择一个子域
        subdomain = random.choice(self.subdomains) if self.subdomains else ""
        
        # 替换占位符
        url = self.tile_url_template
        url = url.replace("{z}", str(zoom))
        url = url.replace("{x}", str(x))
        url = url.replace("{y}", str(y))
        url = url.replace("{s}", subdomain)
        
        return url
    
    def is_ocean_area(self, min_lat, min_lon, max_lat, max_lon):
        """
        检查给定的区域是否主要是海域
        扩展判断规则，包括更多全球海域
        
        参数:
            min_lat, min_lon, max_lat, max_lon: 区域边界经纬度
            
        返回:
            布尔值，表示是否主要是海域
        """
        # 中国东部沿海和东海区域的简单判断
        is_east_china_sea = (min_lon > 120 and max_lon < 130 and 
                             min_lat > 20 and max_lat < 35)
        
        # 黄海区域的简单判断
        is_yellow_sea = (min_lon > 118 and max_lon < 127 and 
                         min_lat > 32 and max_lat < 41)
        
        # 南海区域的简单判断
        is_south_china_sea = (min_lon > 108 and max_lon < 120 and 
                              min_lat > 3 and max_lat < 25)
        
        # 日本海域/太平洋区域
        is_japan_sea = (min_lon > 130 and max_lon < 170 and
                         min_lat > 30 and max_lat < 50)
        
        # 印度洋北部区域
        is_north_indian_ocean = (min_lon > 50 and max_lon < 100 and
                                  min_lat > -10 and max_lat < 25)
        
        # 南印度洋区域
        is_south_indian_ocean = (min_lon > 30 and max_lon < 120 and
                                  min_lat > -50 and max_lat < -10)
        
        return (is_east_china_sea or is_yellow_sea or is_south_china_sea or 
                is_japan_sea or is_north_indian_ocean or is_south_indian_ocean)
    
    def get_alternative_sources(self):
        """
        根据地图类型获取备选瓦片源列表
        
        返回:
            瓦片源列表，按优先级排序
        """
        if self.is_ocean_map:
            # 为海域地图使用特殊的优先级排序
            return [s for s in OCEAN_TILE_PRIORITY if s != self.source_type]
        else:
            # 常规地图使用默认排序
            return [s for s in TILE_SOURCES.keys() if s != self.source_type]
    
    def try_alternative_source(self, zoom, x, y, tried_sources=None):
        """
        尝试使用备选瓦片源下载
        
        参数:
            zoom: 缩放级别
            x: 瓦片X坐标
            y: 瓦片Y坐标
            tried_sources: 已尝试过的瓦片源
            
        返回:
            成功时返回瓦片路径，失败时返回None
        """
        if tried_sources is None:
            tried_sources = [self.source_type]
        
        # 根据地图类型获取备选源
        all_sources = self.get_alternative_sources()
        
        # 选择一个尚未尝试过的瓦片源
        available_sources = [s for s in all_sources if s not in tried_sources]
        
        if not available_sources:
            logger.warning(f"已尝试所有瓦片源，均下载失败 {zoom}/{x}/{y}")
            return None
        
        # 选择下一个可用源
        next_source = available_sources[0]
        
        # 创建一个临时下载器使用备选源
        logger.info(f"尝试备选瓦片源: {next_source} 下载瓦片 {zoom}/{x}/{y}")
        temp_downloader = TileDownloader(
            tile_url_template=TILE_SOURCES[next_source],
            output_dir=self.output_dir,
            max_workers=1,
            retries=1,
            timeout=self.timeout,
            source_type=next_source
        )
        
        # 尝试使用备选源下载
        result = temp_downloader.download_tile(zoom, x, y)
        
        # 更新源的成功率统计
        if next_source not in self.source_success_rates:
            self.source_success_rates[next_source] = {"success": 0, "total": 0}
        
        self.source_success_rates[next_source]["total"] += 1
        if result is not None:
            self.source_success_rates[next_source]["success"] += 1
        
        # 如果备选源也失败，递归尝试下一个源
        if result is None:
            tried_sources.append(next_source)
            return self.try_alternative_source(zoom, x, y, tried_sources)
        
        return result
    
    def get_tile_extension(self):
        """
        根据瓦片源类型获取文件扩展名
        """
        if self.source_type in ["esri", "esri_ocean", "google_satellite"]:
            return ".jpg"  # 这些源通常提供JPG格式
        return ".png"      # 默认使用PNG
    
    def download_tile(self, zoom, x, y):
        """
        下载单个瓦片
        
        参数:
            zoom: 缩放级别
            x: 瓦片X坐标
            y: 瓦片Y坐标
            
        返回:
            成功时返回瓦片路径，失败时返回None
        """
        self.total_requests += 1
        extension = self.get_tile_extension()
        tile_dir = os.path.join(self.output_dir, str(zoom), str(x))
        tile_path = os.path.join(tile_dir, f"{y}{extension}")
        
        # 如果瓦片已存在，跳过下载
        if os.path.exists(tile_path) and os.path.getsize(tile_path) > 100:  # 确保文件不是空的或损坏的
            self.successful_requests += 1
            return tile_path
        
        # 确保瓦片目录存在
        if not os.path.exists(tile_dir):
            os.makedirs(tile_dir)
        
        # 格式化URL
        url = self.format_tile_url(zoom, x, y)
        
        # 添加随机延迟，避免被服务器拒绝
        delay = 0.1 + 0.3 * random.random()
        time.sleep(delay)
        
        # 尝试下载瓦片
        for attempt in range(self.retries):
            try:
                logger.debug(f"尝试下载瓦片 {zoom}/{x}/{y} (尝试 {attempt+1}/{self.retries})")
                
                # 每次请求都使用稍微不同的User-Agent，降低被屏蔽的风险
                if attempt > 0:
                    browser_versions = ["91.0.4472.124", "92.0.4515.107", "93.0.4577.63", "94.0.4606.81"]
                    self.session.headers.update({
                        'User-Agent': f'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/{random.choice(browser_versions)} Safari/537.36'
                    })
                
                response = self.session.get(url, timeout=self.timeout)
                
                if response.status_code == 200:
                    # 检验响应内容，确保不是错误页面或空白图片
                    content_length = len(response.content)
                    if content_length < 100:  # 通常正常的瓦片不会小于100字节
                        logger.warning(f"下载的瓦片 {zoom}/{x}/{y} 内容过小 ({content_length} 字节)，可能是空白图片")
                        time.sleep(1)
                        continue
                    
                    with open(tile_path, 'wb') as f:
                        f.write(response.content)
                    logger.debug(f"成功下载瓦片 {zoom}/{x}/{y}")
                    self.successful_requests += 1
                    return tile_path
                else:
                    logger.warning(f"错误: 状态码 {response.status_code} 下载瓦片 {zoom}/{x}/{y}")
                    # 如果被拒绝，等待更长时间
                    if response.status_code == 403 or response.status_code == 429:
                        wait_time = 2 + attempt * 2  # 指数退避
                        logger.info(f"服务器限制请求，等待 {wait_time} 秒后重试")
                        time.sleep(wait_time)
                    else:
                        time.sleep(1)  # 其他错误稍等一会再重试
            except Exception as e:
                logger.warning(f"下载瓦片 {zoom}/{x}/{y} 时出错: {str(e)}")
                time.sleep(1)  # 失败后稍等一会再重试
        
        self.failed_requests += 1
        
        # 如果当前瓦片源下载失败，尝试备选源
        if self.source_type != "custom":  # 自定义源不尝试替换
            logger.info(f"瓦片 {zoom}/{x}/{y} 使用主源 {self.source_type} 下载失败，尝试备选源")
            return self.try_alternative_source(zoom, x, y)
        
        return None
    
    def download_area(self, min_lat, min_lon, max_lat, max_lon, zoom_levels):
        """
        下载指定区域的瓦片地图
        
        参数:
            min_lat: 最小纬度
            min_lon: 最小经度
            max_lat: 最大纬度
            max_lon: 最大经度
            zoom_levels: 缩放级别列表
            
        返回:
            下载的瓦片数量和失败的瓦片数量
        """
        total_tiles = 0
        failed_tiles = 0
        
        # 检测是否是海域地图
        self.is_ocean_map = self.is_ocean_area(min_lat, min_lon, max_lat, max_lon)
        if self.is_ocean_map:
            logger.info("检测到海域地图，使用优化的海域瓦片源策略")
        
        # 如果是海域且当前使用的不是海域友好的源，建议切换
        if self.is_ocean_map and self.source_type not in ["esri_ocean", "google_satellite", "esri"]:
            logger.warning(f"当前使用的瓦片源 {self.source_type} 对海域支持有限，建议使用 esri_ocean, google_satellite 或 esri")
            
        # 重置统计信息
        self.total_requests = 0
        self.successful_requests = 0
        self.failed_requests = 0
        self.source_success_rates = {}
        
        # 先下载低缩放级别的瓦片
        sorted_zoom_levels = sorted(zoom_levels)
        
        # 提前创建任务以获取总瓦片数
        all_tasks = []
        for zoom in sorted_zoom_levels:
            # 计算边界瓦片坐标
            min_x, max_y = self.deg2num(min_lat, min_lon, zoom)
            max_x, min_y = self.deg2num(max_lat, max_lon, zoom)
            
            # 确保x和y的顺序正确
            min_x, max_x = min(min_x, max_x), max(min_x, max_x)
            min_y, max_y = min(min_y, max_y), max(min_y, max_y)
            
            # 创建下载任务列表
            for x in range(min_x, max_x + 1):
                for y in range(min_y, max_y + 1):
                    all_tasks.append((zoom, x, y))
        
        total_tiles = len(all_tasks)
        if total_tiles == 0:
            logger.warning("计算出的瓦片数量为0，请检查经纬度范围和缩放级别是否正确")
            # 确保至少有一个瓦片用于尝试下载
            logger.info("尝试下载一个测试瓦片以验证下载功能")
            min_zoom = sorted_zoom_levels[0] if sorted_zoom_levels else 1
            test_x, test_y = self.deg2num((min_lat + max_lat) / 2, (min_lon + max_lon) / 2, min_zoom)
            result = self.download_tile(min_zoom, test_x, test_y)
            if result:
                logger.info("测试瓦片下载成功，但计算的瓦片范围为空")
                return 1, 0
            return 0, 0
        
        logger.info(f"总共需要下载 {total_tiles} 个瓦片")
        
        # 按缩放级别分组下载
        for zoom in sorted_zoom_levels:
            # 计算边界瓦片坐标
            min_x, max_y = self.deg2num(min_lat, min_lon, zoom)
            max_x, min_y = self.deg2num(max_lat, max_lon, zoom)
            
            # 确保x和y的顺序正确
            min_x, max_x = min(min_x, max_x), max(min_x, max_x)
            min_y, max_y = min(min_y, max_y), max(min_y, max_y)
            
            # 计算该缩放级别下的瓦片数量
            tiles_count = (max_x - min_x + 1) * (max_y - min_y + 1)
            
            # 如果计算出的瓦片数为0，但经纬度范围有效，尝试添加一个中心点瓦片
            if tiles_count == 0 and min_lat < max_lat and min_lon < max_lon:
                logger.warning(f"缩放级别 {zoom} 计算出的瓦片数为0，尝试添加中心点瓦片")
                center_lat = (min_lat + max_lat) / 2
                center_lon = (min_lon + max_lon) / 2
                center_x, center_y = self.deg2num(center_lat, center_lon, zoom)
                tasks = [(zoom, center_x, center_y)]
                tiles_count = 1
            else:
                logger.info(f"开始下载缩放级别 {zoom} 的瓦片，共 {tiles_count} 个")
                
                # 创建下载任务列表
                tasks = []
                for x in range(min_x, max_x + 1):
                    for y in range(min_y, max_y + 1):
                        tasks.append((zoom, x, y))
            
            # 使用线程池并行下载
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                results = list(tqdm(
                    executor.map(lambda t: self.download_tile(*t), tasks),
                    total=len(tasks),
                    desc=f"下载缩放级别 {zoom} 的瓦片",
                    unit="瓦片"
                ))
            
            # 统计失败的下载
            failed_in_zoom = results.count(None)
            failed_tiles += failed_in_zoom
            
            success_rate = ((len(tasks) - failed_in_zoom) / len(tasks)) * 100 if len(tasks) > 0 else 0
            logger.info(f"缩放级别 {zoom} 已完成: {len(tasks) - failed_in_zoom}/{len(tasks)} 成功 ({success_rate:.2f}%)")
            
            # 如果失败率过高，自动尝试切换到最成功的源
            if failed_in_zoom > 0 and self.source_success_rates:
                best_source = max(self.source_success_rates.items(), 
                                  key=lambda x: x[1]["success"] / x[1]["total"] if x[1]["total"] > 0 else 0)
                best_rate = best_source[1]["success"] / best_source[1]["total"] if best_source[1]["total"] > 0 else 0
                
                if best_rate > 0.7 and best_source[0] != self.source_type:
                    logger.info(f"发现更好的瓦片源: {best_source[0]}，成功率: {best_rate:.2f}，切换源继续下载")
                    self.source_type = best_source[0]
                    self.tile_url_template = TILE_SOURCES[best_source[0]]
                    self.subdomains = SUBDOMAINS.get(best_source[0], ["a"])
            
            # 如果失败率过高，提前结束下载
            if failed_in_zoom > 0 and (len(tasks) - failed_in_zoom) / len(tasks) < 0.1:
                logger.warning(f"缩放级别 {zoom} 下载失败率过高 ({(failed_in_zoom / len(tasks)) * 100:.2f}%)，跳过后续更高缩放级别")
                break
        
        # 确保不会返回0/0瓦片的情况
        if total_tiles == 0 and self.successful_requests > 0:
            total_tiles = self.successful_requests
        
        success_rate = ((total_tiles - failed_tiles) / total_tiles) * 100 if total_tiles > 0 else 0
        logger.info(f"下载完成: {total_tiles - failed_tiles}/{total_tiles} 瓦片下载成功 ({success_rate:.2f}%)")
        
        # 输出各瓦片源的成功率统计
        if self.source_success_rates:
            logger.info("各瓦片源成功率统计:")
            for source, stats in self.source_success_rates.items():
                rate = stats["success"] / stats["total"] * 100 if stats["total"] > 0 else 0
                logger.info(f"  - {source}: {stats['success']}/{stats['total']} ({rate:.2f}%)")
        
        return total_tiles, failed_tiles


if __name__ == "__main__":
    # 示例：下载浙江舟山区域的地图瓦片
    # 舟山大致范围：29.5°N至31°N，121.5°E至123.5°E
    downloader = TileDownloader(output_dir="./tiles", source_type="esri")  # 使用ESRI卫星影像源
    downloader.download_area(
        min_lat=29.5, min_lon=121.5,
        max_lat=31.0, max_lon=123.5,
        zoom_levels=range(5, 13)  # 从5级到12级缩放
    ) 